import lightgbm as lgb
import pandas as pd
import numpy as np
from fastapi import FastAPI
from pydantic import BaseModel
from typing import List, Dict
import os
import threading
import time
# 替换原来的导入
# from pydantic import BaseSettings
from pydantic_settings import BaseSettings

# 创建FastAPI实例
app = FastAPI(title="工知家-CTR预测-LightGBM", version="1.0")

# 模型路径和全局模型变量
model_path = '/data/gongzhijia/data/models/GBDT/final_lgb_model.txt'
current_model = None
new_model = None
model_lock = threading.RLock()  # 用于确保线程安全

# 决策阈值
THRESHOLD = 0.3

# 加载模型函数
def load_model_background():
    global new_model
    try:
        if os.path.exists(model_path):
            # 模拟加载延迟，实际中会由模型大小决定
            # time.sleep(5)  
            new_model = lgb.Booster(model_file=model_path)
            return True
        return False
    except Exception as e:
        print(f"模型加载错误: {str(e)}")
        return False

# 初始加载模型 (同步)
def initial_load_model():
    global current_model
    if os.path.exists(model_path):
        current_model = lgb.Booster(model_file=model_path)
        return True
    return False

# 初始加载模型
initial_load_model()

# 定义请求模型
class PredictionRequest(BaseModel):
    features: List[Dict[str, float]]

# 定义响应模型
class PredictionResponse(BaseModel):
    predictions: List[float]

# 健康检查接口
@app.get("/health")
def health_check():
    return {"status": "healthy", "message": "模型服务正常运行中"}

# 预测接口
@app.post("/predict", response_model=PredictionResponse)
def predict(request: PredictionRequest):
    # 检查模型是否加载
    with model_lock:
        if current_model is None:
            return {"error": "模型未加载成功"}
    
    # 将请求特征转换为DataFrame
    df = pd.DataFrame(request.features)
    
    # 进行预测 (使用当前模型，不阻塞)
    with model_lock:
        predictions = current_model.predict(df)
    
    # 返回结果
    return {
        "predictions": predictions.tolist()
    }

# 重新加载模型接口 - 异步版本
@app.post("/reload-model")
def reload_model():
    try:
        # 启动后台线程加载模型
        thread = threading.Thread(target=load_model_background)
        thread.start()
        
        # 立即返回响应，不等待加载完成
        return {"status": "success", "message": "模型开始在后台加载"}
    except Exception as e:
        return {"status": "error", "message": f"模型加载线程启动失败: {str(e)}"}

# 检查模型加载状态接口
@app.get("/model-status")
def model_status():
    global new_model
    with model_lock:
        if new_model is not None:
            # 模型已加载完成，替换当前模型
            global current_model
            current_model = new_model
            new_model = None  # 重置新模型引用
            return {"status": "ready", "message": "新模型已加载并替换完成"}
        else:
            return {"status": "loading", "message": "模型正在后台加载中"}

# 运行命令: uvicorn main:app --host 0.0.0.0 --port 8000



# 添加配置类
class Settings(BaseSettings):
    model_path: str = '/data/gongzhijia/data/models/GBDT/final_lgb_model.txt'
    threshold: float = 0.3
    port: int = 8001
    host: str = '0.0.0.0'

    class Config:
        env_file = '.env'

# 实例化配置
settings = Settings()

# 修改原来的硬编码
model_path = settings.model_path
THRESHOLD = settings.threshold

# 添加启动代码
if __name__ == '__main__':
    import uvicorn
    uvicorn.run(
        'api-CTR-GBDT:app',
        host=settings.host,
        port=settings.port,
        reload=False,  # 生产环境关闭自动重载
        workers=4  # 根据服务器CPU核心数调整
    )